
import os
import re
from OllamaAPI import OllamaAPI
from prompt import Prompt


class TxtHandle:
    def __init__(self):
        pass
    
    @staticmethod
    def add_biaodianfuhao(string):
        """
        检查句末，添加标点符号
        """
        if string[-1] not in [".", "。", "?", "？", "!","！","”",":","："]:
            string = string + "。"
        return string

    @staticmethod
    def txt_split(txt_path,output_dir,split="^第.+章",encoding="gbk",delete=["【","】"]):
        """
        将txt文件按章节分割
        """
        if not os.path.exists(output_dir):
            print(f"目录不存在，已创建目录: {output_dir}")
            os.makedirs(output_dir)
        with open(txt_path,'r',encoding=encoding) as f1:
            part_name = os.path.join(output_dir,"s001")
            fp = open(part_name,"a+")
            # 获取文件每一行
            for line in f1.readlines():
                line = line.strip()
                if (re.search(split,line)) :
                    for d in delete:
                        line = line.replace(d,"")  # 移除章节号中的特殊字符
                    name = line.replace(" ","")
                    fp.close()
                    part_name = os.path.join(output_dir,name+".txt")
                    fp = open(part_name,"a+")
                    fp.write(line+"\n")
                else:
                    fp.write(line+"\n")

    @staticmethod
    def get_all_txt(path):
        """
        获取目录下所有txt文件
        """
        filepaths = []
        for root, dirs, files in os.walk(path):
            for name in files:
                if '.txt' in name:
                    filepaths.append(os.path.join(root, name))
        return filepaths

    @staticmethod
    def check_empty(input_file,out_file,handle_ban_words=True):
        """
        检查文件是否为空行，为空行则去除
        并检查是否包含违禁词，包含则使用大模型替换*号
        """
        ollama_api = OllamaAPI()

        # 读取文件内容并过滤空白行
        with open(input_file, 'r', encoding='utf-8') as f:
            with open(out_file, 'w', encoding='utf-8') as fw:
                lines = f.readlines()
                # 过滤掉空白行（包括只有空格或制表符的行）
                non_empty_lines = [line for line in lines if line.strip()]
                for index,line in enumerate(non_empty_lines):
                    if "*" in line and handle_ban_words:
                        pre_text = non_empty_lines[index-1]+line+non_empty_lines[index+1]
                        prompt = Prompt.handel_ban_word.format(pre_text,line)
                        options = {"temperature": 0.6, "top_p": 0.9}  # 可选参数
                        respone = ollama_api.generate_text(prompt, options)  # 调用API并获取生成的文本
                        result = respone.split("</think>")[-1]
                        try:
                            result.replace(" ","")
                            text = re.search("\[text:(.+)\]",result,re.DOTALL).group(1)
                        except:
                            fw.write(line+"\n")           # 处理失败，写入原文
                        else:
                            fw.write(text+"\n")             # 写入处理后的内容
                    else:
                        fw.write(line+"\n")
        
        print(f"成功处理文件，输出路径: {out_file}")
        return out_file
    
    @staticmethod
    def delete_target_line(input_file,out_file,target_line:list[str],encoding="utf-8"):
        """
        删除文件中包含目标字符串的行
        """
        with open(input_file, 'r', encoding=encoding) as f:
            with open(out_file, 'w', encoding=encoding) as fw:
                lines = f.readlines()
                for line in lines:
                    line = line.strip()
                    if line == "":
                        continue
                    for target in target_line:
                        if re.search(target,line):
                            break
                    else:
                        fw.write(line+"\n")
        print(f"成功处理文件，输出路径: {out_file}")
        return out_file